package com.luo.ai.assistant.chat.memory;

import com.luo.ai.assistant.util.JsonUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.chat.messages.Message;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.util.CollectionUtils;

import java.util.List;
import java.util.Set;
import java.util.stream.Collectors;

/**
 * Redis存储
 *
 * @author luohq
 * @date 2025-02-12 15:59
 */
public class MyRedisChatMemoryStore {
    private static final Logger log = LoggerFactory.getLogger(MyRedisChatMemoryStore.class);
    /**
     * RedisTemplate
     */
    private RedisTemplate<String, Object> redisTemplate;
    /**
     * Redis ChatMemory配置
     */
    private MyRedisMemoryConfig myRedisMemoryConfig;

    public MyRedisChatMemoryStore(RedisTemplate<String, Object> redisTemplate, MyRedisMemoryConfig myRedisMemoryConfig) {
        this.redisTemplate = redisTemplate;
        this.myRedisMemoryConfig = myRedisMemoryConfig;
    }

    public void addMessage(String conversationId, Long curTimestamp, Message message) {
        log.debug("Redis - add message[{}] at {}: {}", curTimestamp, JsonUtils.toJson(message));
        //保存消息
        this.redisTemplate.opsForZSet()
                .add(this.redisKey(conversationId), message, curTimestamp);
    }

    public List<Message> getMessages(String conversationId, int lastN) {
        //通过Redis ZRANGE命令查询最后的N条聊天记录
        Set<Object> messageSet = this.redisTemplate.opsForZSet()
                .range(this.redisKey(conversationId), -lastN, -1);
        log.debug("Redis - get message history - message count: {}, message list: {}", messageSet.size(), JsonUtils.toJson(messageSet));

        //如果没有历史聊天记录，返回空列表
        if (CollectionUtils.isEmpty(messageSet)) {
            return List.of();
        }
        //转换Set为List
        List<Message> messageList = messageSet.stream()
                .map(Message.class::cast)
                .collect(Collectors.toList());
        return messageList;
    }


    private String redisKey(String conversationId) {
        return this.myRedisMemoryConfig.getKeyPrefix() + conversationId;
    }

}
